from openai import OpenAI
import time

client = OpenAI(
    base_url='http://172.26.63.246:1025/v1/',
    api_key='b011f8fd-99a5-4633-8bf5-fef46cf7bd90', # ModelScope Token
)

response = client.chat.completions.create(
    model='DeepSeek-R1-Distill-Qwen-32B', # ModelScope Model-Id
    messages=[
        {
            'role': 'system',
            'content': '你是一个乐于助人的小助手.'
        },
        {
            'role': 'user',
            'content': '圆周率如何计算'
        }
    ],
    stream=True
)

done_reasoning = False
token_count = 0
start_time = time.time()

for chunk in response:
    # delta.reasoning_content 属性可能不存在
    reasoning_chunk = getattr(chunk.choices[0].delta, 'reasoning_content', '')
    answer_chunk = chunk.choices[0].delta.content

    if reasoning_chunk != '':
        print(reasoning_chunk, end='', flush=True)
        token_count += len(reasoning_chunk.split())  # 假设按空格分隔为 token
    elif answer_chunk != '':
        if not done_reasoning:
            print('\n\n === Final Answer ===\n')
            done_reasoning = True
        print(answer_chunk, end='', flush=True)
        token_count += len(answer_chunk.split())  # 假设按空格分隔为 token

# 计算并打印 token 生成速度
end_time = time.time()
elapsed_time = end_time - start_time
tokens_per_second = token_count / elapsed_time if elapsed_time > 0 else 0

print(f"\n\nToken 生成速度: {tokens_per_second:.2f} tokens/秒")
